import torch
import torch.nn as nn
from torch.optim import SGD
from torch.optim.lr_scheduler import StepLR

# 1. 定义一个极简的线性模型（用于演示）
model = nn.Linear(in_features=10, out_features=1)  # 输入10维，输出1维

# 2. 定义优化器（初始学习率设为0.1）
optimizer = SGD(model.parameters(), lr=0.1)

# 3. 定义StepLR调度器：每5个epoch将学习率乘以0.1（即降为原来的10%）
scheduler = StepLR(optimizer, step_size=5, gamma=0.1)

# 4. 模拟训练循环（共10个epoch）
print("Epoch | 学习率")
print("-" * 15)
for epoch in range(10):
    # 此处省略实际训练步骤（前向/反向传播），仅演示学习率变化
    current_lr = optimizer.param_groups[0]['lr']
    print(f"{epoch+1:^5} | {current_lr:.6f}")
    
    # 调用scheduler.step()更新学习率
    optimizer.step()
    scheduler.step()